{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 4: Federated Learning with Model Averaging\n",
    "\n",
    "**Recap:** In Part 2 of this tutorial, we trained a model using a very simple version of Federated Learning. This required each data owner to trust the model owner to be able to see their gradients.\n",
    "\n",
    "**Description:** In this tutorial, we'll show how to use the advanced aggregation tools from Part 3 to allow the weights to be aggregated by a trusted \"secure worker\" before the final resulting model is sent back to the model owner (us). \n",
    "\n",
    "In this way, only the secure worker can see whose weights came from whom. We might be able to tell which parts of the model changed, but we do NOT know which worker (bob or alice) made which change, which creates a layer of privacy.\n",
    "\n",
    "Authors:\n",
    " - Andrew Trask - Twitter: [@iamtrask](https://twitter.com/iamtrask)\n",
    " - Jason Mancuso - Twitter: [@jvmancuso](https://twitter.com/jvmancuso)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import syft as sy\n",
    "import copy\n",
    "hook = sy.TorchHook(torch)\n",
    "from torch import nn, optim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Create Data Owners\n",
    "\n",
    "First, we're going to create two data owners (Bob and Alice) each with a small amount of data. We're also going to initialize a secure machine called \"secure_worker\". In practice this could be secure hardware (such as Intel's SGX) or simply a trusted intermediary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a couple workers\n",
    "\n",
    "bob = sy.VirtualWorker(hook, id=\"bob\")\n",
    "alice = sy.VirtualWorker(hook, id=\"alice\")\n",
    "secure_worker = sy.VirtualWorker(hook, id=\"secure_worker\")\n",
    "\n",
    "\n",
    "# A Toy Dataset\n",
    "data = torch.tensor([[0,0],[0,1],[1,0],[1,1.]], requires_grad=True)\n",
    "target = torch.tensor([[0],[0],[1],[1.]], requires_grad=True)\n",
    "\n",
    "# get pointers to training data on each worker by\n",
    "# sending some training data to bob and alice\n",
    "bobs_data = data[0:2].send(bob)\n",
    "bobs_target = target[0:2].send(bob)\n",
    "\n",
    "alices_data = data[2:].send(alice)\n",
    "alices_target = target[2:].send(alice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Create Our Model\n",
    "\n",
    "For this example, we're going to train with a simple Linear model. We can initialize it normally using PyTorch's nn.Linear constructor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Iniitalize A Toy Model\n",
    "model = nn.Linear(2,1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Send a Copy of the Model to Alice and Bob\n",
    "\n",
    "Next, we need to send a copy of the current model to Alice and Bob so that they can perform steps of learning on their own datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bobs_model = model.copy().send(bob)\n",
    "alices_model = model.copy().send(alice)\n",
    "\n",
    "bobs_opt = optim.SGD(params=bobs_model.parameters(),lr=0.1)\n",
    "alices_opt = optim.SGD(params=alices_model.parameters(),lr=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Train Bob's and Alice's Models (in parallel)\n",
    "\n",
    "As is conventional with Federated Learning via Secure Averaging, each data owner first trains their model for several iterations locally before the models are averaged together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10):\n",
    "\n",
    "    # Train Bob's Model\n",
    "    bobs_opt.zero_grad()\n",
    "    bobs_pred = bobs_model(bobs_data)\n",
    "    bobs_loss = ((bobs_pred - bobs_target)**2).sum()\n",
    "    bobs_loss.backward()\n",
    "\n",
    "    bobs_opt.step()\n",
    "    bobs_loss = bobs_loss.get().data\n",
    "\n",
    "    # Train Alice's Model\n",
    "    alices_opt.zero_grad()\n",
    "    alices_pred = alices_model(alices_data)\n",
    "    alices_loss = ((alices_pred - alices_target)**2).sum()\n",
    "    alices_loss.backward()\n",
    "\n",
    "    alices_opt.step()\n",
    "    alices_loss = alices_loss.get().data\n",
    "    \n",
    "    print(\"Bob:\" + str(bobs_loss) + \" Alice:\" + str(alices_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 5: Send Both Updated Models to a Secure Worker\n",
    "\n",
    "Now that each data owner has a partially trained model, it's time to average them together in a secure way. We achieve this by instructing Alice and Bob to send their model to the secure (trusted) server. \n",
    "\n",
    "Note that this use of our API means that each model is sent DIRECTLY to the secure_worker. We never see it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alices_model.move(secure_worker)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bobs_model.move(secure_worker)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 6: Average the Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, the last step for this training epoch is to average Bob and Alice's trained models together and then use this to set the values for our global \"model\". "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    model.weight.set_(((alices_model.weight.data + bobs_model.weight.data) / 2).get())\n",
    "    model.bias.set_(((alices_model.bias.data + bobs_model.bias.data) / 2).get())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rinse and Repeat\n",
    "\n",
    "And now we just need to iterate this multiple times!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterations = 10\n",
    "worker_iters = 5\n",
    "\n",
    "for a_iter in range(iterations):\n",
    "    \n",
    "    bobs_model = model.copy().send(bob)\n",
    "    alices_model = model.copy().send(alice)\n",
    "\n",
    "    bobs_opt = optim.SGD(params=bobs_model.parameters(),lr=0.1)\n",
    "    alices_opt = optim.SGD(params=alices_model.parameters(),lr=0.1)\n",
    "\n",
    "    for wi in range(worker_iters):\n",
    "\n",
    "        # Train Bob's Model\n",
    "        bobs_opt.zero_grad()\n",
    "        bobs_pred = bobs_model(bobs_data)\n",
    "        bobs_loss = ((bobs_pred - bobs_target)**2).sum()\n",
    "        bobs_loss.backward()\n",
    "\n",
    "        bobs_opt.step()\n",
    "        bobs_loss = bobs_loss.get().data\n",
    "\n",
    "        # Train Alice's Model\n",
    "        alices_opt.zero_grad()\n",
    "        alices_pred = alices_model(alices_data)\n",
    "        alices_loss = ((alices_pred - alices_target)**2).sum()\n",
    "        alices_loss.backward()\n",
    "\n",
    "        alices_opt.step()\n",
    "        alices_loss = alices_loss.get().data\n",
    "    \n",
    "    alices_model.move(secure_worker)\n",
    "    bobs_model.move(secure_worker)\n",
    "    with torch.no_grad():\n",
    "        model.weight.set_(((alices_model.weight.data + bobs_model.weight.data) / 2).get())\n",
    "        model.bias.set_(((alices_model.bias.data + bobs_model.bias.data) / 2).get())\n",
    "    \n",
    "    print(\"Bob:\" + str(bobs_loss) + \" Alice:\" + str(alices_loss))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lastly, we want to make sure that our resulting model learned correctly, so we'll evaluate it on a test dataset.  In this toy problem, we'll use the original data, but in practice we'll want to use new data to understand how well the model generalizes to unseen examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "preds = model(data)\n",
    "loss = ((preds - target) ** 2).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(preds)\n",
    "print(target)\n",
    "print(loss.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this toy example, the averaged model is underfitting relative to a plaintext model trained locally would behave, however we were able to train it without exposing each worker's training data.  We were also able to aggregate the updated models from each worker on a trusted aggregator to prevent data leakage to the model owner.\n",
    "\n",
    "In a future tutorial, we'll aim to do our trusted aggregation directly with the gradients, so that we can update the model with better gradient estimates and arrive at a stronger model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Congratulations!!! - Time to Join the Community!\n",
    "\n",
    "Congratulations on completing this notebook tutorial! If you enjoyed this and would like to join the movement toward privacy preserving, decentralized ownership of AI and the AI supply chain (data), you can do so in the following ways!\n",
    "\n",
    "### Star PySyft on GitHub\n",
    "\n",
    "The easiest way to help our community is just by starring the Repos! This helps raise awareness of the cool tools we're building.\n",
    "\n",
    "- [Star PySyft](https://github.com/OpenMined/PySyft)\n",
    "\n",
    "### Join our Slack!\n",
    "\n",
    "The best way to keep up to date on the latest advancements is to join our community! You can do so by filling out the form at [http://slack.openmined.org](http://slack.openmined.org)\n",
    "\n",
    "### Join a Code Project!\n",
    "\n",
    "The best way to contribute to our community is to become a code contributor! At any time you can go to PySyft GitHub Issues page and filter for \"Projects\". This will show you all the top level Tickets giving an overview of what projects you can join! If you don't want to join a project, but you would like to do a bit of coding, you can also look for more \"one off\" mini-projects by searching for GitHub issues marked \"good first issue\".\n",
    "\n",
    "- [PySyft Projects](https://github.com/OpenMined/PySyft/issues?q=is%3Aopen+is%3Aissue+label%3AProject)\n",
    "- [Good First Issue Tickets](https://github.com/OpenMined/PySyft/issues?q=is%3Aopen+is%3Aissue+label%3A%22good+first+issue%22)\n",
    "\n",
    "### Donate\n",
    "\n",
    "If you don't have time to contribute to our codebase, but would still like to lend support, you can also become a Backer on our Open Collective. All donations go toward our web hosting and other community expenses such as hackathons and meetups!\n",
    "\n",
    "[OpenMined's Open Collective Page](https://opencollective.com/openmined)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
